#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from std_msgs.msg import Float64
from geometry_msgs.msg import PoseStamped
import tkinter as tk
from tkinter import ttk
import threading

class AltitudeController(Node):
    """
    高度控制节点
    提供简单的GUI界面设置目标高度，并发布到/target_altitude话题
    """
    
    def __init__(self):
        super().__init__('altitude_controller')
        
        # 创建发布者
        self.altitude_pub = self.create_publisher(
            Float64, 
            '/target_altitude', 
            10
        )
        
        # 创建订阅者，监听当前位置
        self.pose_sub = self.create_subscription(
            PoseStamped,
            '/drone_pose',  # 假设这是当前位置话题
            self.pose_callback,
            10
        )
        
        # 当前状态
        self.current_altitude = 0.0
        self.target_altitude = 1.0  # 默认目标高度1米
        
        # 定时发布目标高度
        self.timer = self.create_timer(0.1, self.publish_target_altitude)
        
        # 启动GUI（在单独线程中）
        self.gui_thread = threading.Thread(target=self.create_gui, daemon=True)
        self.gui_thread.start()
        
        self.get_logger().info('高度控制节点已启动')
        
    def pose_callback(self, msg):
        """更新当前高度"""
        self.current_altitude = msg.pose.position.z
        
    def publish_target_altitude(self):
        """发布目标高度"""
        msg = Float64()
        msg.data = self.target_altitude
        self.altitude_pub.publish(msg)
        
    def create_gui(self):
        """创建GUI界面"""
        self.root = tk.Tk()
        self.root.title("无人机高度控制")
        self.root.geometry("300x200")
        
        # 当前高度显示
        current_frame = ttk.Frame(self.root)
        current_frame.pack(pady=10)
        
        ttk.Label(current_frame, text="当前高度:").pack(side=tk.LEFT)
        self.current_label = ttk.Label(current_frame, text="0.00 m")
        self.current_label.pack(side=tk.LEFT, padx=10)
        
        # 目标高度设置
        target_frame = ttk.Frame(self.root)
        target_frame.pack(pady=10)
        
        ttk.Label(target_frame, text="目标高度:").pack(side=tk.LEFT)
        self.altitude_var = tk.StringVar(value="1.00")
        self.altitude_entry = ttk.Entry(
            target_frame, 
            textvariable=self.altitude_var,
            width=10
        )
        self.altitude_entry.pack(side=tk.LEFT, padx=5)
        ttk.Label(target_frame, text="m").pack(side=tk.LEFT)
        
        # 设置按钮
        button_frame = ttk.Frame(self.root)
        button_frame.pack(pady=10)
        
        set_button = ttk.Button(
            button_frame,
            text="设置高度",
            command=self.set_altitude
        )
        set_button.pack(side=tk.LEFT, padx=5)
        
        # 快速设置按钮
        quick_frame = ttk.Frame(self.root)
        quick_frame.pack(pady=10)
        
        ttk.Label(quick_frame, text="快速设置:").pack()
        quick_buttons = ttk.Frame(quick_frame)
        quick_buttons.pack()
        
        for height in [0.5, 1.0, 1.5, 2.0]:
            btn = ttk.Button(
                quick_buttons,
                text=f"{height}m",
                command=lambda h=height: self.quick_set_altitude(h)
            )
            btn.pack(side=tk.LEFT, padx=2)
        
        # 定时更新显示
        self.update_display()
        
        self.root.mainloop()
        
    def set_altitude(self):
        """设置目标高度"""
        try:
            new_altitude = float(self.altitude_var.get())
            if 0.0 <= new_altitude <= 10.0:  # 限制高度范围
                self.target_altitude = new_altitude
                self.get_logger().info(f'目标高度设置为: {new_altitude:.2f}m')
            else:
                self.get_logger().warn('高度超出范围 (0-10m)')
        except ValueError:
            self.get_logger().error('无效的高度值')
            
    def quick_set_altitude(self, height):
        """快速设置高度"""
        self.altitude_var.set(f"{height:.2f}")
        self.set_altitude()
        
    def update_display(self):
        """更新显示"""
        if hasattr(self, 'current_label'):
            self.current_label.config(text=f"{self.current_altitude:.2f} m")
        
        # 每100ms更新一次
        if hasattr(self, 'root'):
            self.root.after(100, self.update_display)

def main(args=None):
    rclpy.init(args=args)
    
    node = AltitudeController()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main() 